'''
Author: Zhuangyu
Date: 2022-06-02 14:31:11
LastEditTime: 2022-07-05 11:34:18
'''
from __future__ import print_function, division
import time
#Importing relevant modules and packages
import numpy as np
import matplotlib.pyplot as plt
import random
from Parameters import *
from scipy import integrate, linalg, interpolate
from tqdm import tqdm
from casadi import *
import mpctools as mpc
import casadi
import scipy.io as sio  
import van_Genuchten as vg

#Space and time parameters
Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
DeltaT,Nsim=time_parameters()
p=Loam() #Soil parameters

Na=Nz
Np=Na+Nz
Nu=1
C=SX.eye(Nz)
C1=SX.zeros((1,Nz))
C1[0]=1
size=Nsim
Ny1=1

def RichardsModel_1D(x,u,ui):
    
    #x is the state vector
    #u is the irrigation amount (the input)
    #w is the process noise
    #Kc is the crop coefficient
    #ET0 is the reference evapotranspiration
    #p represents the soil parameters
    
    Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
    # dhdt=SX.zeros(Nz)
    dhdt=SX.zeros(Nz)
    
    Kc=0.88      #Crop coefficient

    ET0=1.389e-08 #Reference evapotranspiration
    
    for i in range(0,Nz):

        CurrentState=x[i]         
                
        if i == 0: #Bottom boundary, Free drainage boundary condition(BC)
            BCz_L=CurrentState
            BCz_U=x[i+1]
         
        elif i == Nz-1: #Top boundary
            BCz_L = x[i - 1]    
         
        else: #Internal nodes
            BCz_L = x[i - 1]
            BCz_U = x[i + 1]

        #Computing the unsaturated hydraulic conductivity at the centre of the compartments
        KzL=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_L,p))
        KzU=0.5*(vg.KFun(CurrentState,p)+vg.KFun(BCz_U,p))
        
        #Computing the pressure head gradient
        DHzL=(CurrentState-BCz_L)/dz 
        DHzU=(BCz_U-CurrentState)/dz

        if i == Nz-1: #Surface Node [Nuemann BC which incorporates the irrigation amount, u]
            Term1 = (1.0/dz)*( -u - KzL*DHzL)
            Term2 =(1.0/dz)*( -1*KzL)
        
        else: #Nodes below the surface
            Term1 = (1.0/dz) * (KzU*DHzU - KzL*DHzL)
            Term2 = (1.0/dz) * (KzU - KzL)

        Term3 = (Kc*ET0)/Hz #The Sink term [Root water extraction rate]
        Term4 = Term1 + Term2 - Term3
        Term5 = Term4/vg.CFun(CurrentState,p)

        dhdt[i]=Term5
    DHDT=dhdt + ui
    return DHDT


x=SX.sym('x',Nz)
u=SX.sym('u',Nu)
ui=SX.sym('ui',Na)

F_symbol_c = RichardsModel_1D(x,u,ui)
ODE_c = {'x': x, 'p': vertcat(u,ui), 'ode': F_symbol_c}
opts = {'tf': DeltaT, 'regularity_check':True}  # seconds
I = integrator('I', 'cvodes', ODE_c, opts)  # Build casadi integrator


def getOutputs1D_allnodes(x):
    #Returns the measurements for all the states
    Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,Hz=circular_parameters()
    y = vg.thetaFun(x,p)
    CMatrix = cmatrix_single(Nz)
    # y_pred = mtimes(CMatrix,y)
    y_pred = mtimes(CMatrix,y)
    return y_pred

def getOutputs1D(C,x):
    #Returns measurements for some selected states based on the C-matrix
    y = vg.thetaFun(x,p)
    # y_pred = mtimes(C,y)
    y_pred = mtimes(C,y)    
    return y_pred

def pre_jacobian(x,u,ui):
    
    F = RichardsModel_1D(x,u,ui)
    F_jacx=jacobian(F,x)
    F_jacu=jacobian(F,u)
    F_jacui=jacobian(F,ui)
    
    return F_jacx,F_jacu,F_jacui

def autoNorm(data):

    mins = data.min(0)
    maxs = data.max(0)
    ranges = -(maxs-mins)
    normData = np.zeros(np.shape(data))
    row = data.shape[0]
    normData = data-np.tile(mins,(row,1))
    normData = normData/np.tile(ranges,(row,1))
    return normData

F_symbol1=pre_jacobian(x,u,ui)[0]
F_symbol2=pre_jacobian(x,u,ui)[1]
F_symbol3=pre_jacobian(x,u,ui)[2]


F1=Function('F1',[x,u,ui],[F_symbol1])
F2=Function('F2',[x,u,ui],[F_symbol2])
F3=Function('F3',[x,u,ui],[F_symbol3])

def state_sensitivity(state_sen, jac_state, jac_pars):
    dhdt = mtimes(jac_state, state_sen) + jac_pars
    return dhdt

def state_sensitivity_np(state_sen, jac_state, jac_pars):
    dhdt = np.matmul(jac_state, state_sen) + jac_pars
    return dhdt

#Add one step that integrates the sensisitivity equation with respect to the state
def state_sensitivity_approx(state_sen, jac_state, jac_pars):
    xx_sen_cur = state_sen
    for i in range(12):
        xx_sen_cur = xx_sen_cur + 60*state_sensitivity_np(xx_sen_cur, jac_state, jac_pars) # 1 min sampling time is used
        pass
    return xx_sen_cur


def ode_sensitivity(x,u,ui, x_sensitivity):
    
    E1=time.time() 
    dfdx_matrix = F1(x,u,ui)
    
    #print("dfdx matrix --- Start------")
    #print(dfdx_matrix.full().max())
    #print(dfdx_matrix.full().min())
    #print("dfdx matrix --- End------")
    
    dfdp_matrix = F3(x,u,ui)
    #print("dfdp matrix --- Start------")
    #print(dfdp_matrix.full().max())
    #print(dfdp_matrix.full().min())
    #print("dfdp matrix --- End ------")
    dfdx_init = np.zeros((Nz, Nz)) # The state equation does not explictly depend on the initial state
    #dfdx0_matrix = dfdx_matrix[:,Nx:]
    #df_x_init = np.zeros((Nx_aug, Nx))
    dfdx_act = dfdx_init
    #xk_new = state_sensitivity_approx(x_sensitivity, dfdx_matrix, dfdx0_matrix)
    xk_new = state_sensitivity_approx(x_sensitivity, dfdx_matrix, dfdx_act)
    
    #print("Sensitivity matrix --- Start------")
    #print(xk_new.full().min())
    #print(xk_new.full().max())
    #print("Sensitivity matrix --- End ------")
    E2=time.time()-E1
    print('Sensitivity ODE takes',E2,' seconds to compute')
    return xk_new




u=np.zeros(6*120)
uu=np.zeros(Nsim)
for i in range(6*120):
    if i < 8:
        u[i]=-5e-06 #Moisture is supplied only at the first time instant
    else:
        u[i]=0
uu=np.tile(u,(1,60))
uu=np.transpose(uu)


#Initial condition [pressure head]
x0=-1.5*np.ones(Nz)

#Numpy array to store the simulation results [pressure head]
xx=np.zeros((Nsim+1,Nz))
xx[0,:]=x0

ode1=np.zeros((Nsim+1,Nz))

#Model uncertainty
np.random.seed(2)
w=0.00001*np.random.randn(Nsim,Nz)
# w=0.01*np.random.randn(Nsim,Nz)
np.random.seed(3)
vv=0.01*np.random.randn(Nsim+1,Nz) #measurement noise

#Numpy array to store the simulation results [volumetric moisture content]
yy=np.zeros((Nsim+1,Nz))
yy[0,:]=getOutputs1D_allnodes(x0).full().ravel()+vv[0,:]

x0_sensitivity = np.zeros((Nz, Nz)) # The system is not augmented


for i in range(Nz):
    x0_sensitivity[i, i] = 1
    pass
xx_sensitivity = x0_sensitivity

#unknown inputs
a1=0*np.ones((Nsim+1,Nz))
a=0.0000005*np.ones((Nsim+1,Nz))

e0 = time.time()

for i in tqdm(range(1, Nsim+1)):
    # Generate observed data (x bias)
    Ik1 = I(x0=xx[i-1], p=vertcat(uu[i-1],a[i-1]))
    xk1 = Ik1['xf'] .full().ravel()
    xx[i,:]=xk1+w[i-1]
    yy[i,:]=getOutputs1D_allnodes(xx[i,:]).full().ravel()+vv[i,:]
    
# define the H jacobian matrix
H=getOutputs1D(C,x)
h_jacx=jacobian(H,x)
Fy = Function('Fy', [x], [h_jacx])

sens_matrix = np.zeros((Ny1*(Nsim+1), Nz))
sens_matrix_Norm = np.zeros((Ny1*(Nsim+1), Nz))

#dhdx_matrix = Fy(x0_aug)
dhdx_matrix = Fy(x0)
dhdx_matrix = dhdx_matrix[1,:]
dhdp_matrix = np.zeros((1,Na))
dhdx_init = np.zeros((Ny1, Nz)) # The output equation does not explcitly depend on the initial state

dhdp=dhdx_init
sens_matrix0 = mtimes(dhdx_matrix, xx_sensitivity)
sens_matrix0 = sens_matrix0 + dhdp

sens_matrix[0:Ny1, :] = sens_matrix0.full()
sens_matrix0_Norm = sens_matrix0.full()

for j in range(sens_matrix0_Norm.shape[0]):
    sens_matrix0_Norm[j, :] = sens_matrix0_Norm[j, :] / yy[0, j]
    pass

sens_matrix_Norm[0:Ny1, :] = abs(sens_matrix0_Norm)

for i in range(1, Nsim+1):
#for i in range(1, 1+1):
    print(i)    
    d=i-1
    e1=time.time()
    
    #Simulate the nonlinear model 
    Ik1 = I(x0=xx[i-1], p=vertcat(uu[i-1],a[i-1]))
    xk1 = Ik1['xf'] .full().ravel()
    xx[i,:]=xk1+w[i-1]
    yy[i,:]=getOutputs1D_allnodes(xx[i,:]).full().ravel()+vv[i,:]
    
    #Compute the jacobians for the output sensitivity
    #Calculate the sensitivity with respect to the state
    xx_sensitivity = ode_sensitivity(xx[i,:], uu[i-1],a[i], xx_sensitivity)
    dhdx_matrix = Fy(xx[i,:])
    dhdx_matrix = dhdx_matrix[1,:]
    dhdp_matrix = np.zeros((1,Na))
    dhdx_init = np.zeros((Ny1, Nz))
    dhdp = dhdx_init
    #dhdx0_matrix = dhdx_matrix[:,Nx:]
    #dhdx_init = np.zeros((Ny, Nx))
    sens_matrix0 = mtimes(dhdx_matrix, xx_sensitivity)
    sens_matrix0 = sens_matrix0 + dhdp
    #print(sens_matrix0.max())
    #print(sens_matrix0.min())
    #print(xx_sensitivity.max())
    #print(xx_sensitivity.min())
    #sens_matrix0 += dhdx0_matrix
    #sens_matrix0 += dhdx_act
    sens_matrix[(i)*Ny1:(i+1)*Ny1,:] = sens_matrix0.full()
    sens_matrix0_Norm = sens_matrix0.full()
    print(sens_matrix0.full().max())
    print(sens_matrix0.full().min())
    
    # normalize the sensitivity matrix. Since the model already scaled parameters, here we dont need to do it again
    for j in range(sens_matrix0_Norm.shape[0]):
        sens_matrix0_Norm[j,:] = sens_matrix0_Norm[j,:]/yy[i,j]

    sens_matrix_Norm[(i)*Ny1:(i+1)*Ny1,:] = abs(sens_matrix0_Norm)

    sum_inTime = np.sum(sens_matrix_Norm, axis=0)
    
    e2=time.time()-e1
    print('The original model simulation takes',e2,'seconds to solve')
    
elapsed_time=time.time()-e0
print(elapsed_time)



rank = np.linalg.matrix_rank(sens_matrix_Norm)
rank_S = rank
cond = np.linalg.cond(sens_matrix_Norm)
print("rank:",rank)
print("cond:",cond)


# t= list(range(Nsim+1))
# # plt.figure(1)
# # for i in range (Nz):
# #     plt.subplot(5, 7, i+1)
# #     plt.plot(t,xx[:,i],'r',t,xx1[:,i],'b')
    

# plt.figure(2)
# for i in range (Nz):
#     plt.subplot(4, 4, i+1)
#     plt.plot(t,xx[:,i],'r')
# plt.show()